﻿module AssemblyCodeGenerator
open EmitType
open Environment
open StateMonad
open Ast
open PrintUtils

let rec private assembly (emit  : EmitType.Emit) : State<unit, Environment> =
  state{
    match emit with
    | EmitType.Emit.Expr (expr, dest_t) ->
      return! assembly_expr expr dest_t
    | EmitType.Emit.Stmt (stmt) ->
      return! assembly_stmt stmt
    | Emit.Comment str ->
      do! comment str
      
    | Emit.Label label ->
      do! print_label !label

    | EmitType.Emit.Jmp (jmp) ->
      return! assembly_jmp jmp
  }

and assembly_expr expr dest_t : State<unit, Environment> =
  state{
    let! dest_t_add = get_t_address dest_t    
    let! dest_t_dword = get_t_DWORD_address dest_t
    let reduce op t0 t1 =
        state{
          let! t0_dword = get_t_DWORD_address t0
          let! t0 = get_t_address t0
          let! t1 = get_t_address t1
          do! mov MemReg.EAX (MemOrReg t0)
          do! op MemReg.EAX (MemOrReg t1)
          do! mov dest_t_dword MemRegConst.EAX
        }

    match expr with
    | EmitExpr.Add (t0, t1) -> 
      do! reduce add t0 t1      

    | EmitExpr.Sub (t0, t1) -> 
      do! reduce sub t0 t1

    | EmitExpr.And (t0, t1) -> 
      do! reduce _and t0 t1

    | EmitExpr.Or (t0, t1) -> 
      do! reduce _or t0 t1

    | EmitExpr.Div (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0      
      let! t0 = get_t_address t0
      let! t1 = get_t_address t1
      do! mov MemReg.EAX (MemOrReg t0)
      do! cdq //this will clear EDX due to the sign extension
      do! idiv MemReg.EAX (MemOrReg t1)
      do! mov dest_t_dword MemRegConst.EAX

    | EmitExpr.Mul (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0
      let! t0 = get_t_address t0
      let! t1 = get_t_address t1
      do! mov MemReg.EAX (MemOrReg t0)
      do! imul eax (MemOrReg t1)
      do! mov dest_t_dword MemRegConst.EAX

    | EmitExpr.Mod (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0      
      let! t0 = get_t_address t0
      let! t1 = get_t_address t1
      do! mov MemReg.EAX (MemOrReg t0)
      do! cdq //this will clear EDX due to the sign extension
      do! idiv MemReg.EAX (MemOrReg t1)
      do! mov dest_t_dword MemRegConst.EDX

    | EmitExpr.Not t -> 
      let! t_dword = get_t_DWORD_address t
      do! neg eax
      do! sbb eax MemRegConst.EAX
      do! add MemReg.EAX (MemRegConst.ConstOfInt 1)
      do! movzx MemReg.EAX MemRegConst.AL
      do! mov dest_t_dword MemRegConst.EAX

    | EmitExpr.Geq (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0
      do! mov MemReg.EAX (MemOrReg t0_dword)
      do! set_zero MemReg.ECX
      let! t1_dword = get_t_DWORD_address t1
      do! cmp MemReg.EAX (MemOrReg t1_dword)
      do! set "ge"
      do! mov dest_t_dword MemRegConst.ECX
    | EmitExpr.Gt (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0
      do! mov MemReg.EAX (MemOrReg t0_dword)
      do! set_zero MemReg.ECX
      let! t1_dword = get_t_DWORD_address t1
      do! cmp MemReg.EAX (MemOrReg t1_dword)
      do! set "g"
      do! mov dest_t_dword MemRegConst.ECX
    | EmitExpr.Leq (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0
      do! mov MemReg.EAX (MemOrReg t0_dword)
      do! set_zero MemReg.ECX
      let! t1_dword = get_t_DWORD_address t1
      do! cmp MemReg.EAX (MemOrReg t1_dword)
      do! set "le"
      do! mov dest_t_dword MemRegConst.ECX
    | EmitExpr.Lt (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0
      do! mov MemReg.EAX (MemOrReg t0_dword)
      do! set_zero MemReg.ECX
      let! t1_dword = get_t_DWORD_address t1
      do! cmp MemReg.EAX (MemOrReg t1_dword)
      do! set "l"
      do! mov dest_t_dword MemRegConst.ECX
    | EmitExpr.Neq (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0
      do! mov MemReg.EAX (MemOrReg t0_dword)
      do! set_zero MemReg.ECX
      let! t1_dword = get_t_DWORD_address t1
      do! cmp MemReg.EAX (MemOrReg t1_dword)
      do! set "ne"
      do! mov dest_t_dword MemRegConst.ECX    
    | EmitExpr.Eq (t0, t1) -> 
      let! t0_dword = get_t_DWORD_address t0
      do! mov MemReg.EAX (MemOrReg t0_dword)
      do! set_zero MemReg.ECX
      let! t1_dword = get_t_DWORD_address t1
      do! cmp MemReg.EAX (MemOrReg t1_dword)
      do! set "e"
      do! mov dest_t_dword MemRegConst.ECX
          
    | EmitExpr.Neg (t) -> 
      do! set_zero MemReg.EAX
      let! t_dword = get_t_DWORD_address t
//      do! cmp t_dword (MemRegConst.ConstOfInt(0))
//      do! set "e"
      do! sub MemReg.EAX (MemOrReg t_dword)
      do! mov dest_t_dword MemRegConst.EAX
      (*do! neg t_dword
      do! mov MemReg.EBX (MemOrReg t_dword)
      do! mov dest_t_dword MemRegConst.EBX*)

    | EmitExpr.ReadArray (t0, t1) -> 
      let! t0 = get_t_address t0
      do! mov MemReg.EBX (MemRegConst.MemOrReg t0)
      let! t1 = get_t_address t1
      do! mov MemReg.ECX (MemRegConst.MemOrReg t1)
      do! imul ecx (MemRegConst.ConstOfInt(4))
//      let! array_value = sum_registers_with_mult_offset ebx ecx Const.Four
      do! add MemReg.EBX MemRegConst.ECX
      do! mov MemReg.EBX (MemOrReg (PtrReg(ebx, 0)))
      do! mov dest_t_dword (MemRegConst.EBX)

    | EmitExpr.GetField (t, id) -> 
      let! t = get_t_address t
      do! mov MemReg.EAX (MemOrReg t)
      do! mov MemReg.EBX (MemOrReg(MemReg.PtrReg(eax, (id + 1) * 4)))
      do! mov dest_t_dword MemRegConst.EBX //(MemOrReg(MemReg.PtrReg(eax, id)))

    | EmitExpr.DirectCall (name, t_list) -> 
      let t_list_inv = t_list |> List.rev
      for t in t_list_inv do
        let! t = get_t_address t
        do! push (MemOrReg t)

      do! call name 
      do! add MemReg.ESP (MemRegConst.ConstOfInt(t_list.Length * 4))
      do! mov dest_t_dword MemRegConst.EAX

    | EmitExpr.Call (id, t_list) -> 
      let t_list_inv = t_list |> List.rev
      for t in t_list_inv do
        let! t = get_t_address t
        do! push (MemOrReg t)
      let! obj = get_t_address t_list.Head
      do! mov MemReg.EAX (MemOrReg obj)
      do! mov MemReg.EAX (MemOrReg (PtrReg(eax, 0)))
      do! call ("[%eax+" + (string (id * 4)) + "]")
      do! add MemReg.ESP (MemRegConst.ConstOfInt(t_list.Length * 4))
      do! mov dest_t_dword MemRegConst.EAX

    | EmitExpr.Int _int -> 
      do! mov dest_t_dword (MemRegConst.ConstOfInt _int)

    | EmitExpr.Str _string -> 
      let! str = add_const _string
      do! lea ebx (Mem.Mem str)
      do! mov dest_t_dword MemRegConst.EBX

    | EmitExpr.Bool _bool -> 
      do! mov dest_t_dword (MemRegConst.ConstOfBool _bool)

    | EmitExpr.Id t -> 
      let! t = get_t_address t 
      do! mov MemReg.EAX (MemOrReg t)
      do! mov dest_t_dword MemRegConst.EAX

    | EmitExpr.Length t -> 
      let! t_add = get_t_address t
      do! mov MemReg.EAX (MemOrReg t_add)
      let! dreg = (sum_registers_with_offset eax -4)
      do! mov MemReg.EBX (MemOrReg dreg)
      do! mov dest_t_dword MemRegConst.EBX //(MemOrReg dreg)

    | EmitExpr.New ty -> 
      let! state = getState
      let dimension = state.GlobalSymbols.[ty].ObjectLayout.Count
      do! push (MemRegConst.ConstOfInt(dimension * 4))
      do! call "__allocateObject"
      do! add MemReg.ESP (MemRegConst.ConstOfInt(4))
      do! lea ebx (Mem.Mem("_" + ty + "_DV"))
      do! mov MemReg.EAX.GetRegAddress MemRegConst.EBX
      do! mov dest_t_add MemRegConst.EAX

    | EmitExpr.NewArray (ty, t) -> 
      let! t_add = get_t_address t
      do! mov MemReg.EAX (MemOrReg t_add)
      do! mov MemReg.EBX MemRegConst.EAX
      do! add MemReg.EAX (MemRegConst.ConstOfInt(1))
      do! imul eax (MemRegConst.ConstOfInt 4)
      do! push (MemRegConst.EAX)
      do! call "__allocateObject"
      do! add MemReg.ESP (MemRegConst.ConstOfInt(4))
      do! mov MemReg.EAX.GetRegAddress MemRegConst.EBX
      do! add MemReg.EAX (MemRegConst.ConstOfInt(4))
      do! mov dest_t_add MemRegConst.EAX

    | EmitExpr.This -> 
      let! this = get_t_this
      do! mov MemReg.EAX (MemOrReg this)
      do! mov dest_t_add (MemOrReg MemReg.EAX)
    
    | EmitExpr.Null -> 
      do! mov dest_t_dword (MemRegConst.ConstOfInt 0)

  }

and assembly_stmt stmt =
  state{
    match stmt with
    | EmitStmt.Assign (t_var, t_value) ->
      let! var_add = get_t_DWORD_address t_var
      let! val_add = get_t_DWORD_address t_value
      do! mov MemReg.EAX (MemOrReg val_add)
      do! mov var_add MemRegConst.EAX
    | EmitStmt.ReturnVoid ->
      do! leave
      do! ret
    | EmitStmt.Return t_ret ->
      let! val_add = get_t_DWORD_address t_ret
      do! mov MemReg.EAX (MemOrReg val_add)
      do! leave
      do! ret
    | EmitStmt.SetArray (t_arr, t_offset, t_val) -> 
      let! t0 = get_t_address t_arr
      do! mov MemReg.EBX (MemRegConst.MemOrReg t0)
      let! t1 = get_t_address t_offset
      do! mov MemReg.ECX (MemRegConst.MemOrReg t1)
      do! imul ecx (MemRegConst.ConstOfInt(4))
      do! add MemReg.EBX MemRegConst.ECX
//      do! mov MemReg.EBX (MemOrReg (PtrReg(ebx, 0)))
//      let! array_value = sum_registers_with_mult_offset ebx ecx Const.Four
      let! value_address = get_t_address t_val
      do! mov MemReg.EAX (MemOrReg value_address)
      do! mov (PtrReg (ebx, 0)) MemRegConst.EAX
    | EmitStmt.SetField (t_var, id, t_val) -> 
      let! t_var = get_t_address t_var
      let! t_val = get_t_address t_val
      do! mov MemReg.EAX (MemOrReg t_var)
      do! mov MemReg.EBX (MemOrReg t_val)
      do! mov (MemReg.PtrReg(eax, (id + 1) * 4)) MemRegConst.EBX
  }

and assembly_jmp _jmp = 
  state{
    match _jmp with
    | EmitJmp.Tjmp (t0, label) ->
      let! cond_add = get_t_address t0
      do! mov MemReg.EAX (MemOrReg cond_add)
      do! test MemReg.EAX MemRegConst.EAX
      do! je !label
    | EmitJmp.Fjmp (t0, label) ->
      let! cond_add = get_t_address t0
      do! mov MemReg.EAX (MemOrReg cond_add)
      do! test MemReg.EAX MemRegConst.EAX
      do! jne !label
    | EmitJmp.Jmp label ->
      do! jmp !label
  }

and generate_code (code : Emit list) =
  state{
    
    match code with
    | [] -> 
      return ()
    | line :: lines ->
      do! assembly line
      return! generate_code lines
  }

let private generate_prologue print_code c_name name = 
  do print_code ""
  do print_code "#--------------------------------------------------"
  do print_code "    .align 4"
  do print_code ""
  do print_code (sprintf "_%s_%s:" c_name name)
  do print_code ""
  do print_code "    push     %ebp                #  prologue"
  do print_code "    mov     %ebp, %esp"

let private generate_epilogue print_code c_name name = 
  do print_code ""
  do print_code (sprintf "_epilogue_%s_%s:" c_name name) 
  do print_code "    leave # epilogue"
  do print_code "    ret"
  do print_code ""

let private generate_method_code c_name state (name : string, args : T_symbol list, stack : T_symbol list, code : Emit list) = 
  let print_code = state.print_string_nl
  do generate_prologue print_code c_name name
  do state.ActualArgs <- args.Length
  if args.Length > 0 then
    do print_code ("    sub     %esp, " + (string (stack.Length * 4)))
  do print_code "    push     %ebx"
  do print_code ""
  for i, t in stack |> Seq.mapi (fun i t -> i + 1, t) do
    do print_code ("    mov     DWORD PTR [%ebp+-" + (string (i * 4)) + "], 0      #  initialize locals")
  do print_code ""
      
  do execute (generate_code code) (state) |> ignore
  do generate_epilogue print_code c_name name

let private print_dispatch_vectors (global_symbols : SymbolTableTypes.GlobalSymbols) (print) = 
  for c in global_symbols.Values do
    do print ""
    do print (sprintf "_%s_DV:" c.Name) 
    for name, c_name in c.DispatchVector do
      print (sprintf "    .long _%s_%s" !c_name name)
    do print ""

let private print_entry_point (global_symbols : SymbolTableTypes.GlobalSymbols) (print) = 
  let main = 
    let maybe_ret = global_symbols.Values |> Seq.tryFind (fun c -> c.DispatchVector |> Seq.exists (fun (name, _) -> name = "main"))
    match maybe_ret with
    | Some res ->
      res
    | _ ->
      failwith "Error: no main in program"

  do generate_prologue print "_tj" "main"

  do print (sprintf "    push %d        #  o = new C" (4 * (main.ObjectLayout.Count + 1)))
  do print "    call     __allocateObject         #  %eax -> o"
  do print "    add %esp, 4"
  do print ("    lea     %ebx, " + (sprintf "_%s_DV" main.Name) + "             #  init DV reference")
  do print "    mov     [%eax], %ebx"

  do print "    push     DWORD PTR [%ebp+8]      #  o.main[args]"
  do print "    push     %eax                     #  push this"
  do print "    mov     %eax, DWORD PTR [%eax]"
  do print ("    call     [%eax+" + (main.DispatchVector |> Seq.findIndex (fun (s, _) -> s = "main") |> (*) 4 |> string) + "]")
  do print "    add     %esp, 8"

  do print "    mov     %eax, 0                #  return 0"

  do generate_epilogue print "ic" "main"

let private print_exit_points print = 
  do print "# ----------------------------------------------------"
  do print ""
  do print "    .align 4"
  do print (!EmitType.NULL_POINTER + ":")
  do print "    lea     %eax, strNPC"
  do print "    push     %eax"
  do print "    call     __println"
  do print "    push     1"
  do print "    call     __exit"

  do print ""

  do print "    .align 4"
  do print (!EmitType.ARRAY_OUT_OF_BOUNDS + ":")
  do print "    lea     %eax, strABC"
  do print "    push     %eax"
  do print "    call     __println"
  do print "    push     1"
  do print "    call     __exit"

let private print_arch print prg_name = 
  do print "  #  global declarations"
  do print (sprintf "    .file     \"%s\"" prg_name)
  do print "    .intel_syntax"
  do print ""
  do print "#  data section"
  do print ".section .rodata"

  do print "    .align 4"

let print_code_section print = 
  do print "#  text (code) section"
  do print ".section .text"
  do print ""
  do print "    .globl    __tj_main"
  do print ""

let generate_main_code prg_name symbols (emit : (EmitType.ClassName * EmitType.ClassFields * EmitType.ClassMethods) list) dest = 

  let consts : ResizeArray<string> = ResizeArray()
  let code : ResizeArray<string> = ResizeArray()

  let print_code (s:string) = code.Add s
  let print_const (s:string) = consts.Add s
  
//  do System.IO.File.WriteAllLines("", code)
  let state = 
    {
      GlobalSymbols     = symbols
      NumTabs           = 2
      Consts            = ResizeArray()
      print_string_nl   = print_code
      ActualArgs        = 0
    }

  do print_code_section print_code

  for (c_name, fields, ms) in emit do
    for m in ms do
      do generate_method_code c_name state m 

  do print_entry_point state.GlobalSymbols print_code

  do print_exit_points print_code

  do print_arch print_const prg_name

  do print_dispatch_vectors state.GlobalSymbols print_const

  do print_const "    .int 23"
  do print_const "strNPC:"
  do print_const "    .string \"Null pointer violation.\""
  do print_const "    .int 23"
  do print_const "strABC:     "
  do print_const "    .string \"Array bounds violation.\""
  
  do for i,c in state.Consts |> Seq.mapi (fun i c -> i, c)  do
      do print_const ""
      do print_const (sprintf "    .int %d" c.Length)
      do print_const (sprintf "str%d:" i)
      do print_const (sprintf "    .string \"%s\"" c)
      do print_const ""
    
  //let path = "D:\Università\Compilers\out.s"
//  let path = "out.s"
  let c_path = System.IO.Path.GetFullPath(dest)
  
  do System.IO.File.WriteAllLines(dest, Seq.append consts code)

